{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "PyTorch on Cloud TPUs: Single Core Training AlexNet on Fashion MNIST",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "TPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IZgybIMXbpCl",
        "colab_type": "text"
      },
      "source": [
        "##PyTorch on Cloud TPUs: Single Core Training AlexNet on Fashion MNIST \n",
        "\n",
        "This notebook trains the [AlexNet](https://arxiv.org/abs/1404.5997) network on the [Fashion MNIST dataset](https://github.com/zalandoresearch/fashion-mnist) using PyTorch with a single Cloud TPU core. This will show you how to train your own networks on a single Cloud TPU core and highlight the differences between using one vs. many Cloud TPU cores.\n",
        "\n",
        "PyTorch can use Cloud TPU cores as devices with the PyTorch/XLA package. For more on PyTorch/XLA see its [Github](https://github.com/pytorch/xla) or its [documentation](http://pytorch.org/xla/). We also have a [\"Getting Started\"](https://colab.research.google.com/github/pytorch/xla/blob/master/contrib/colab/getting-started.ipynb) Colab notebook. Additional Colab notebooks, like this one, are available on the PyTorch/XLA Github.\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qxSMwPAIb5zI",
        "colab_type": "text"
      },
      "source": [
        "### Installing PyTorch/XLA\n",
        "\n",
        "Run the following cell (or copy it into your own notebook!) to install PyTorch, Torchvision, and PyTorch/XLA. It will take a couple minutes to run."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tdIxkkiACX92",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install cloud-tpu-client==0.10 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.9-cp37-cp37m-linux_x86_64.whl"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WyREUv4CCYHH",
        "colab_type": "text"
      },
      "source": [
        "### Only run the below commented cell if you would like a nightly release"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "20zLDvraY2mk",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# VERSION = \"20200325\"  #@param [\"1.5\" , \"20200325\", \"nightly\"]\n",
        "# !curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n",
        "# !python pytorch-xla-env-setup.py --version $VERSION"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u_izjOUccDg4",
        "colab_type": "text"
      },
      "source": [
        "### Dataset & Network\n",
        "\n",
        "In this notebook we'll train AlexNet on the Fashion MNIST dataset. Both are provided by the [Torchvision package](https://pytorch.org/docs/stable/torchvision/index.html).\n",
        "\n",
        "Before diving in, let's look at the Fashion MNIST dataset. The dataset has 10 classes, each represented by an integer index. The following cell creates a mapping from these indices to their corresponding human-readable strings, then downloads the Fashion MNIST training dataset from Torchvision. It may take a minute to run.\n",
        "\n",
        "Torchvision provides easy access to many datasets, including COCO, CIFAR, and Cityscapes. See its documentation for a complete list. \n",
        "\n",
        "The dataset is stored on the Colab VM in the `/tmp/fashionmnist` directory.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-FkLA4MSlcNf",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Downloads the Fashion MNIST dataset using Torchvision\n",
        "# Note: This may take a minute.\n",
        "import os\n",
        "import torch\n",
        "import torchvision\n",
        "import torchvision.datasets as datasets\n",
        "\n",
        "# Maps integer class index to human-readable string describing the class\n",
        "class_map = {\n",
        "0 : \"t-shirt\",\n",
        "1 : \"trouser\",\n",
        "2 : \"pullover\",\n",
        "3 : \"dress\",\n",
        "4 : \"coat\",\n",
        "5 : \"sandal\",\n",
        "6 : \"shirt\",\n",
        "7 : \"sneaker\",\n",
        "8 : \"bag\",\n",
        "9 : \"ankle boot\"\n",
        "}\n",
        "\n",
        "# Downloads the Fashion MNIST dataset using Torchvision\n",
        "raw_dataset = datasets.FashionMNIST(\n",
        "  os.path.join(\"/tmp/fashionmnist\"),\n",
        "  train=True,\n",
        "  download=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zRyIvDK6Z_xM",
        "colab_type": "text"
      },
      "source": [
        "Now that we have the dataset, we can look at elements of it directly. Each element of raw_dataset is a tuple. The first element is a single channel greyscale 28x28 Python Image Library (PIL) image, and the second element is the example's class index. You can change the `img_index` in the following cell to visualize different examples."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "40bWSY2_ZrO9",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "img_index = 0\n",
        "tup = raw_dataset[img_index]\n",
        "display(tup[0].resize((224, 224)))\n",
        "print(class_map[tup[1]])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4BzG5t8HbaMf",
        "colab_type": "text"
      },
      "source": [
        "Now we'll create our AlexNet model and put it on a Cloud TPU core.\n",
        "\n",
        "Torchvision lets us download many model architectures, including VGG, ResNet, and MobileNet. See its documentation for a complete list. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xMMZxHoIk9n7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import torch_xla\n",
        "import torch_xla.core.xla_model as xm\n",
        "\n",
        "\n",
        "# Creates AlexNet for 10 classes\n",
        "net = torchvision.models.alexnet(num_classes=10)\n",
        "\n",
        "# Acquires the default Cloud TPU core and moves the model to it\n",
        "device = xm.xla_device()\n",
        "net = net.to(device)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RYB7dKkqfsiI",
        "colab_type": "text"
      },
      "source": [
        "### Dataloaders\n",
        "\n",
        "Now that we have our dataset and network, let's look at how we'll transform and load the data. As we saw when looking at the Fashion MNIST dataset (above), the examples are 28x28 single channel greyscale PIL images. Torchvision networks, however, expect PyTorch tensors representing  normalized three channel RGB images that are at least 224x224. The following cell defines a conversion from the original examples to the PyTorch tensors our AlexNet requires.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AoTe3v5W881g",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import torchvision.transforms as transforms\n",
        "\n",
        "# See https://pytorch.org/docs/stable/torchvision/models.html for normalization\n",
        "# Pre-trained TorchVision models expect RGB (3 x H x W) images\n",
        "# H and W should be >= 224\n",
        "# Loaded into [0, 1] and normalized as follows:\n",
        "normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
        "                                 std=[0.229, 0.224, 0.225])\n",
        "to_rgb = transforms.Lambda(lambda image: image.convert('RGB'))\n",
        "resize = transforms.Resize((224, 224))\n",
        "my_transform = transforms.Compose([resize, to_rgb, transforms.ToTensor(), normalize])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tFTMqbNMgwWI",
        "colab_type": "text"
      },
      "source": [
        "Now we can download train and test datasets using TorchVision and apply this transform to them."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gsXFl-CFgtG9",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "train_dataset = datasets.FashionMNIST(\n",
        "  os.path.join(\"/tmp/fashionmnist\"),\n",
        "  train=True,\n",
        "  download=True,\n",
        "  transform=my_transform)\n",
        "\n",
        "test_dataset = datasets.FashionMNIST(\n",
        "  os.path.join(\"/tmp/fashionmnist\"),\n",
        "  train=False,\n",
        "  download=True,\n",
        "  transform=my_transform)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YV9nMJRSh_ld",
        "colab_type": "text"
      },
      "source": [
        "PyTorch provides a variety of \"Samplers\" to acquire elements of datasets in different orders. The [Random Sampler](https://pytorch.org/docs/stable/data.html#torch.utils.data.RandomSampler), which we'll use, samples elements randomly without replacement."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "khRAFXTe-KfX",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "train_sampler = torch.utils.data.RandomSampler(train_dataset)\n",
        "test_sampler = torch.utils.data.RandomSampler(test_dataset)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dy1Eq2O_iagC",
        "colab_type": "text"
      },
      "source": [
        "Finally, we create DataLoaders that load batches of examples from our Dataset according to our Sampler's policy. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TV5yqz2W9PzX",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "batch_size = 8\n",
        "\n",
        "train_loader = torch.utils.data.DataLoader(\n",
        "  train_dataset,\n",
        "  batch_size=batch_size,\n",
        "  sampler=train_sampler)\n",
        "\n",
        "test_loader = torch.utils.data.DataLoader(\n",
        "  test_dataset,\n",
        "  batch_size=batch_size,\n",
        "  sampler=test_sampler)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kHQbrK5ijXGS",
        "colab_type": "text"
      },
      "source": [
        "### Evaluating the Network\n",
        "\n",
        "Before training we want to verify our data->network pipeline and set a baseline level of performance. The following cell defines a function, `eval_network,` that runs a network on a dataset and shows the percentage of the dataset that was correctly classified as well as a sample batch."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UlSLfLaW_A3G",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import time\n",
        "from google.colab import widgets\n",
        "\n",
        "t_to_img = transforms.Compose([transforms.ToPILImage(), transforms.Grayscale()])\n",
        "\n",
        "# Runs the given net on the batches provided by the test_loader\n",
        "# Records the number of correct predictions (guesses) and \n",
        "# prints the percentage of correct guesses on the dataset, plus a \n",
        "# sample batch.\n",
        "def eval_network(net, test_loader):\n",
        "  start_time = time.time()\n",
        "  num_correct = 0\n",
        "  total_guesses = 0\n",
        "\n",
        "  # Sets eval and no grad context for evaluation\n",
        "  net.eval()\n",
        "  with torch.no_grad():\n",
        "    for data, targets in iter(test_loader):\n",
        "      # Sends data and targets to device\n",
        "      data = data.to(device)\n",
        "      targets = targets.to(device)\n",
        "\n",
        "      # Acquires the network's best guesses at each class\n",
        "      results = net(data)\n",
        "      best_guesses = torch.argmax(results, 1)\n",
        "\n",
        "      # Updates number of correct and total guesses\n",
        "      num_correct += torch.eq(targets, best_guesses).sum().item()\n",
        "      total_guesses += batch_size\n",
        "    \n",
        "    # Prints output\n",
        "    elapsed_time = time.time() - start_time\n",
        "    print(\"Correctly guessed \", num_correct/total_guesses*100, \"% of the dataset\")\n",
        "    print(\"Evaluated in \", elapsed_time, \" seconds\")\n",
        "    print(\"Sample batch:\")\n",
        "    \n",
        "    # Uses last batch as sample\n",
        "    grid = widgets.Grid(2, 4)\n",
        "    row = 0\n",
        "    col = 0\n",
        "    for ex in zip(data, targets, best_guesses):\n",
        "      data = ex[0].cpu()\n",
        "      target = class_map[ex[1].item()]\n",
        "      guess = class_map[ex[2].item()]\n",
        "\n",
        "      img = t_to_img(data)\n",
        "\n",
        "      with grid.output_to(row, col):\n",
        "        display(img)\n",
        "        print(\"Target: \", target)\n",
        "        print(\"Guess: \", guess)\n",
        "\n",
        "        # Updates grid location\n",
        "        if col == 3:\n",
        "          row += 1\n",
        "          col = 0\n",
        "        else:\n",
        "          col += 1"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EyXjocsog0We",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Percentage of guesses that are correct (expected to be 0.1)\n",
        "eval_network(net, test_loader)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qencz3bJpaNU",
        "colab_type": "text"
      },
      "source": [
        "You should see the untrained network guess about 10% of the dataset correctly, since it's randomly guessing and there are 10 classes."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oJFsnpfxf0j2",
        "colab_type": "text"
      },
      "source": [
        "### Training the Network\n",
        "\n",
        "Randomly guessing at what we're seeing isn't very interesting, so it's time to train our network.\n",
        "\n",
        "Basic training in PyTorch:\n",
        "\n",
        "- enumerates a dataset in batches\n",
        "- runs the network on each batch\n",
        "- evaluates the network's performance using a loss function\n",
        "- calls `backward()` to propagate gradients through the network\n",
        "- uses an optimizer to `step()` and apply the gradients to the network's weights\n",
        "\n",
        "In this case we'll use the [CrossEntropyLoss](https://pytorch.org/docs/master/nn.html#torch.nn.CrossEntropyLoss) since we have a classification problem and AlexNet returns unnormalized scores for each class. We'll also use the [Adam](https://pytorch.org/docs/master/optim.html#torch.optim.Adam) optimizer since it's a popular optimizer.\n",
        "\n",
        "We'll only train for a single epoch for timeliness. The following cell should take 4-5 minutes to run."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RSl2WMQMGutS",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Note: this will take 5-10 minutes to run.\n",
        "num_epochs = 1\n",
        "loss_fn = torch.nn.CrossEntropyLoss()\n",
        "optimizer = torch.optim.Adam(net.parameters())\n",
        "\n",
        "# Ensures network is in train mode\n",
        "net.train()\n",
        "\n",
        "start_time = time.time()\n",
        "for epoch in range(num_epochs):\n",
        "  for data, targets in iter(train_loader):\n",
        "    # Sends data and targets to device\n",
        "    data = data.to(device)\n",
        "    targets = targets.to(device)\n",
        "\n",
        "    # Acquires the network's best guesses at each class\n",
        "    results = net(data)\n",
        "\n",
        "    # Computes loss\n",
        "    loss = loss_fn(results, targets)\n",
        "\n",
        "    # Updates model\n",
        "    optimizer.zero_grad()\n",
        "    loss.backward()\n",
        "    xm.optimizer_step(optimizer, barrier=True)  # Note: Cloud TPU-specific code!\n",
        "\n",
        "elapsed_time = time.time() - start_time\n",
        "print (\"Spent \", elapsed_time, \" seconds training for \", num_epochs, \" epoch(s) on a single core.\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "id8KQMotr3DA",
        "colab_type": "text"
      },
      "source": [
        "\n",
        "So far we've just used Cloud TPU cores as devices, like any other PyTorch device. In the above training loop we see our first Cloud TPU-specific line of code: `xm.optimizer_step(optimizer, barrier=True).` \n",
        "\n",
        "PyTorch uses Cloud TPUs through the [XLA deep learning compiler](https://www.tensorflow.org/xla). This compiler records the operations we perform into a graph which is evaluated all-at-once and as needed. `xm.optimizer_step(optimizer, barrier=True)` inserts a \"barrier\" in the graph that forces evaluation every time the gradients are updated. This prevents XLA's graphs from growing too large. For more details about how PyTorch uses XLA and Cloud TPUs see [the documentation](http://pytorch.org/xla/).\n",
        "\n",
        "\n",
        "We can now re-evaluate our network's performance:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KA_74FcEZ5cq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "eval_network(net, test_loader)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YHSM9PGcsdfd",
        "colab_type": "text"
      },
      "source": [
        "You should see the network dramatically improve (it should guess 60%+)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tl8iYcWxgJUv",
        "colab_type": "text"
      },
      "source": [
        "### What's Next\n",
        "\n",
        "This Colab showed how to train AlexNet on the Fashion MNIST dataset using a single Cloud TPU core. Using a Cloud TPU core as a device required one tweak to our network: instead of calling `optimizer.step()` we called `optimizer_step(optimizer, barrier=True).` In future Colabs we'll see a few other tweaks that will let us maximize PyTorch's performance on Cloud TPUs. The [documentation](http://pytorch.org/xla/) also describes them. \n",
        "\n",
        "Additional information about PyTorch on Cloud TPUs, including more sample Colabs, is available on the PyTorch/XLA package's [Github](https://github.com/pytorch/xla). You're encouraged to try PyTorch on Cloud TPUs on both Colab and Google Cloud Platform, too! For Colab, just copy the first code cell in this notebook to start your own. If you have ideas/suggestions/comments the best way to contact the team is with an [issue on our Github](https://github.com/pytorch/xla/issues). "
      ]
    }
  ]
}
